import numpy as np
from common.img_utils.sas_python.utils import (
    vector_idx_minus_1_from_matlab as vector_idx_minus_1_to_start_0,
)

from skimage.future.graph import RAG
import math


def rag_mean_color_revised(image, labels, connectivity=1, mode="distance", sigma=255.0):
    """
    REVISED FROM THE ORIGINAL VERSION (skimage.future.graph.rag_mean_color)
    Compute the Region Adjacency Graph using mean colors.

    Given an image and its initial segmentation, this method constructs the
    corresponding Region Adjacency Graph (RAG). Each node in the RAG
    represents a set of pixels within `image` with the same label in `labels`.
    The weight between two adjacent regions represents how similar or
    dissimilar two regions are depending on the `mode` parameter.

    Parameters
    ----------
    image : ndarray, shape(M, N, [..., P,] 3)
        Input image.
    labels : ndarray, shape(M, N, [..., P])
        The labelled image. This should have one dimension less than
        `image`. If `image` has dimensions `(M, N, 3)` `labels` should have
        dimensions `(M, N)`.
    connectivity : int, optional
        Pixels with a squared distance less than `connectivity` from each other
        are considered adjacent. It can range from 1 to `labels.ndim`. Its
        behavior is the same as `connectivity` parameter in
        ``scipy.ndimage.generate_binary_structure``.
    mode : {'distance', 'similarity'}, optional
        The strategy to assign edge weights.

            'distance' : The weight between two adjacent regions is the
            :math:`|c_1 - c_2|`, where :math:`c_1` and :math:`c_2` are the mean
            colors of the two regions. It represents the Euclidean distance in
            their average color.

            'similarity' : The weight between two adjacent is
            :math:`e^{-d^2/sigma}` where :math:`d=|c_1 - c_2|`, where
            :math:`c_1` and :math:`c_2` are the mean colors of the two regions.
            It represents how similar two regions are.
    sigma : float, optional
        Used for computation when `mode` is "similarity". It governs how
        close to each other two colors should be, for their corresponding edge
        weight to be significant. A very large value of `sigma` could make
        any two colors behave as though they were similar.

    Returns
    -------
    out : RAG
        The region adjacency graph.

    Examples
    --------
    >>> from skimage import data, segmentation
    >>> from skimage.future import graph
    >>> img = data.astronaut()
    >>> labels = segmentation.slic(img)
    >>> rag = graph.rag_mean_color(img, labels)

    References
    ----------
    .. [1] Alain Tremeau and Philippe Colantoni
           "Regions Adjacency Graph Applied To Color Image Segmentation"
           :DOI:`10.1109/83.841950`
    """
    graph = RAG(labels, connectivity=connectivity)
    if len(graph.nodes) == 0:
        return None

    for n in graph:
        graph.nodes[n].update(
            {
                "labels": [n],
                "pixel count": 0,
                "pixel index": [],
                "total color": np.array([0] * image.shape[2], dtype=np.double),
                "total var color": np.array([0] * image.shape[2], dtype=np.double),
            }
        )

    abs_indexes = np.array(list(range(labels.shape[0] * labels.shape[1])))
    abs_indexes_img = np.reshape(abs_indexes, labels.shape[::-1]).T

    for index in np.ndindex(labels.shape):
        current = labels[index]
        abs_index = abs_indexes_img[index]
        graph.nodes[current]["pixel count"] += 1
        graph.nodes[current]["pixel index"].append(abs_index)
        graph.nodes[current]["total color"] += image[index]

    for n in graph:
        graph.nodes[n]["mean color"] = (
            graph.nodes[n]["total color"] / graph.nodes[n]["pixel count"]
        )
        graph.nodes[n]["pixel index"] = np.sort(np.array(graph.nodes[n]["pixel index"]))

    # Calculate Variance of colors in a node
    for index in np.ndindex(labels.shape):
        current = labels[index]
        graph.nodes[current]["total var color"] += (
            image[index] - graph.nodes[current]["mean color"]
        ) ** 2
    for n in graph:
        graph.nodes[n]["var color"] = (
            graph.nodes[n]["total var color"] / graph.nodes[n]["pixel count"]
        )

    for x, y, d in graph.edges(data=True):
        diff = graph.nodes[x]["mean color"] - graph.nodes[y]["mean color"]
        diff = np.linalg.norm(diff)
        if mode == "similarity":
            d["weight"] = math.e ** (-(diff ** 2) / sigma)
        elif mode == "distance":
            d["weight"] = diff
        else:
            raise ValueError("The mode '%s' is not recognised" % mode)

    return graph


def calculate_superpixel_props(
    img, label_img, connectivity=1, mode="distance", sigma=255.0
):
    """Calculate the post-superpixel results for pre-processing in the SAS algorithm.
    :param img:
    :param label_img:
    :param connectivity: 1 means 4-connectivity; 2 means 8-connectivity
    :param mode: copied from rag_mean_color_revised
    :param sigma: copied from rag_mean_color_revised
    :return:

    Parameters
    ----------
    img : ndarray
        ndarray image
    label_img : numpy array
        image label map (h x w)
    connectivity : int, optional
        [description], by default 1
    mode : str, optional
        [description], by default 'distance'
    sigma : float, optional
        [description], by default 255.0

    Returns
    -------
    seg : dict
        A dictionary which has the length of (# of superpixel).
        In each key, the values are pixels' indices which belong to this superpixel.
    seg_vals : list
        A list which has the length of (# of superpixel).
        In each element, the values are this superpixel's mean BGR(N) values or RGB(N) values.
    seg_edges: list
        A ndarray include all superpixels' edges ((# of edge pairs) x 2)
    seg_vars : list
        A list which has the length of (# of superpixel).
        In each element, the values are this superpixel's variance of BGR(N) values or RGB(N) values.
    """
    label_img = vector_idx_minus_1_to_start_0(label_img)
    g = rag_mean_color_revised(
        img, label_img, connectivity=connectivity, mode=mode, sigma=sigma
    )

    if g is None:
        return {}, [], [], []

    seg = {}
    seg_vals = []
    seg_vars = []
    for n in g.nodes:
        seg[str(n)] = g.nodes[n]["pixel index"]
        seg_vals.append(g.nodes[n]["mean color"])
        seg_vars.append(g.nodes[n]["var color"])
    seg_vals = np.array(seg_vals)
    seg_vars = np.array(seg_vars)

    g_directed = g.to_directed()
    seg_edges = np.array(g_directed.edges)

    return seg, seg_vals, seg_edges, seg_vars
